import os
import subprocess
import tempfile
import time
import json
import re  # 用于解析数组格式的输出
from typing import Dict, List, Optional, Tuple, Any

class Judge:
    def __init__(self, time_limit: float = 2.0, memory_limit: int = 256):
        self.time_limit = time_limit
        self.memory_limit = memory_limit
        self.lang_ext = {
            'cpp': '.cpp',
            'python': '.py',
            'java': '.java'
        }
        # 模板根目录（指向 code_templates 文件夹）
        self.templates_root = os.path.join(
            os.path.dirname(os.path.dirname(__file__)), 
            'code_templates'
        )

    def _load_test_cases_from_files(self, testcases_dir: str) -> List[Dict[str, str]]:
        """从目录读取测试用例（.in和.out文件对）"""
        test_cases: List[Dict[str, str]] = []
        if not os.path.exists(testcases_dir):
            print(f"测试用例目录不存在: {testcases_dir}")
            return test_cases

        for filename in sorted(os.listdir(testcases_dir)):
            if filename.endswith(".in"):
                base_name = filename[:-3]
                input_path = os.path.join(testcases_dir, filename)
                output_path = os.path.join(testcases_dir, f"{base_name}.out")

                if os.path.exists(output_path):
                    try:
                        with open(input_path, 'r', encoding='utf-8') as f_in:
                            input_data = f_in.read().strip()
                        with open(output_path, 'r', encoding='utf-8') as f_out:
                            expected_output = f_out.read().strip()

                        test_cases.append({
                            "input_data": input_data,
                            "expected_output": expected_output
                        })
                    except Exception as e:
                        print(f"读取测试用例文件失败 ({input_path}): {str(e)}")

        return test_cases

    def render_full_code(self, core_code: str, language: str, function_signature: str, template_file: str) -> Tuple[str, Optional[str]]:
        """拼接核心代码与题目对应的模板文件，生成完整可运行代码"""
        template_path = os.path.join(self.templates_root, template_file)
        
        if not os.path.exists(template_path):
            return "", f"模板文件不存在: {template_path}"

        try:
            with open(template_path, 'r', encoding='utf-8') as f:
                template = f.read()
        except Exception as e:
            return "", f"读取模板文件失败: {str(e)}"

        # 替换模板中的占位符
        full_code = template.replace("{{FUNCTION_SIGNATURE}}", function_signature)\
                           .replace("{{CORE_CODE}}", core_code)
        return full_code, None

    def run_code(self, full_code: str, language: str, input_data: str) -> Tuple[str, str, float, str]:
        """编译/运行代码，返回状态、输出、时间、错误"""
        with tempfile.TemporaryDirectory() as temp_dir:
            if language == "cpp":
                return self._run_cpp(full_code, temp_dir, input_data)
            elif language == "python":
                return self._run_python(full_code, temp_dir, input_data)
            elif language == "java":
                return self._run_java(full_code, temp_dir, input_data)
            else:
                return "unsupported_language", "", 0.0, f"不支持的语言: {language}"

    def _run_cpp(self, full_code: str, temp_dir: str, input_data: str) -> Tuple[str, str, float, str]:
        """运行C++代码"""
        cpp_path = os.path.join(temp_dir, "solution.cpp")
        exe_path = os.path.join(temp_dir, "solution")
        
        with open(cpp_path, 'w', encoding='utf-8') as f:
            f.write(full_code)

        # 编译
        compile_cmd = ["g++", cpp_path, "-o", exe_path, "-std=c++11", "-Wall"]
        compile_result = subprocess.run(
            compile_cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True
        )

        if compile_result.returncode != 0:
            error_msg = f"编译错误:\n{compile_result.stderr}"
            return "compilation_error", "", 0.0, error_msg

        # 运行
        try:
            start_time = time.time()
            run_result = subprocess.run(
                [exe_path],
                input=input_data,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
                timeout=self.time_limit
            )
            end_time = time.time()
            execution_time = end_time - start_time

            if run_result.returncode != 0:
                error_msg = f"运行时错误:\n{run_result.stderr}"
                return "runtime_error", "", execution_time, error_msg

            output = run_result.stdout.strip()
            return "accepted", output, execution_time, ""

        except subprocess.TimeoutExpired:
            return "time_limit_exceeded", "", self.time_limit, f"超时（超过{self.time_limit}秒）"
        except Exception as e:
            return "runtime_error", "", 0.0, f"运行错误: {str(e)}"

    def _run_python(self, full_code: str, temp_dir: str, input_data: str) -> Tuple[str, str, float, str]:
        """运行Python代码"""
        py_path = os.path.join(temp_dir, "solution.py")
        
        with open(py_path, 'w', encoding='utf-8') as f:
            f.write(full_code)

        try:
            start_time = time.time()
            result = subprocess.run(
                ['python3', py_path],
                input=input_data.encode(),
                capture_output=True,
                timeout=self.time_limit
            )
            end_time = time.time()

            execution_time = end_time - start_time

            if result.returncode != 0:
                return "runtime_error", "", execution_time, result.stderr.decode('utf-8')

            actual_output = result.stdout.decode('utf-8').strip()
            return "accepted", actual_output, execution_time, ""
        except subprocess.TimeoutExpired:
            return "time_limit_exceeded", "", self.time_limit, "执行时间超过限制"
        except Exception as e:
            return "runtime_error", "", 0.0, f"运行错误: {str(e)}"

    def _run_java(self, full_code: str, temp_dir: str, input_data: str) -> Tuple[str, str, float, str]:
        """运行Java代码"""
        java_path = os.path.join(temp_dir, "Main.java")
        
        with open(java_path, 'w', encoding='utf-8') as f:
            f.write(full_code)

        # 编译
        compile_cmd = ["javac", java_path]
        compile_result = subprocess.run(
            compile_cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True
        )

        if compile_result.returncode != 0:
            error_msg = f"编译错误:\n{compile_result.stderr}"
            return "compilation_error", "", 0.0, error_msg

        # 运行
        try:
            start_time = time.time()
            run_result = subprocess.run(
                ["java", "-cp", temp_dir, "Main"],
                input=input_data,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
                timeout=self.time_limit
            )
            end_time = time.time()
            execution_time = end_time - start_time

            if run_result.returncode != 0:
                error_msg = f"运行时错误:\n{run_result.stderr}"
                return "runtime_error", "", execution_time, error_msg

            output = run_result.stdout.strip()
            return "accepted", output, execution_time, ""

        except subprocess.TimeoutExpired:
            return "time_limit_exceeded", "", self.time_limit, f"超时（超过{self.time_limit}秒）"
        except Exception as e:
            return "runtime_error", "", 0.0, f"运行错误: {str(e)}"

    def _parse_and_sort_output(self, output_str: str) -> Optional[str]:
        """解析数组格式输出（如"[1,2]"），排序后返回标准化字符串"""
        try:
            # 去除所有空格，处理 "[1, 2]" 或 "[1,2]" 等格式
            output_clean = re.sub(r'\s+', '', output_str)
            # 检查是否为数组格式
            if not (output_clean.startswith('[') and output_clean.endswith(']')):
                return None  # 非数组格式不处理
            
            # 提取元素（空数组直接返回）
            elements = output_clean[1:-1].split(',')
            if not elements[0]:  # 处理 "[]" 空数组
                return "[]"
            
            # 转换为数字（支持整数和浮点数）并排序
            elements_num = []
            for elem in elements:
                if '.' in elem:
                    elements_num.append(float(elem))
                else:
                    elements_num.append(int(elem))
            elements_num.sort()  # 排序
            
            # 转回字符串格式
            return '[' + ','.join(map(str, elements_num)) + ']'
        except Exception as e:
            print(f"解析输出失败: {output_str}, 错误: {e}")
            return None  # 解析失败则不处理

    def judge_submission(self, core_code: str, language: str, problem: Any) -> Dict[str, Any]:
        """判题主函数：支持有序/无序输出比对"""
        results: List[Dict[str, Any]] = []
        overall_status = "accepted"
        total_time = 0.0

        # 1. 获取模板文件和函数签名
        template_file = problem.template_file
        if not template_file:
            return {
                "status": "error",
                "results": [],
                "execution_time": 0,
                "error": "题目未配置模板文件"
            }

        func_sig = problem.function_signature.get(language, "")
        if not func_sig:
            return {
                "status": "error",
                "results": [],
                "execution_time": 0,
                "error": "缺少函数签名定义"
            }

        # 2. 生成完整代码
        full_code, error = self.render_full_code(
            core_code, language, func_sig, template_file
        )
        if error:
            return {
                "status": "error",
                "results": [],
                "execution_time": 0,
                "error": error
            }

        # 3. 加载测试用例
        test_cases = self._load_test_cases_from_files(problem.testcases_dir)
        if not test_cases:
            return {
                "status": "error",
                "results": [],
                "execution_time": 0,
                "error": f"未找到测试用例文件: {problem.testcases_dir}"
            }

        # 4. 处理每个测试用例（核心：根据output_type处理比对逻辑）
        for test_case in test_cases:
            run_status, output, execution_time, error = self.run_code(
                full_code, language, test_case["input_data"]
            )

            # 若运行成功，检查输出是否匹配
            if run_status == "accepted":
                # 获取题目输出类型（默认有序）
                output_type = getattr(problem, "output_type", "ordered")
                actual_output = output.strip()
                expected_output = test_case["expected_output"].strip()

                # 对无序输出题目进行标准化处理
                if output_type == "unordered":
                    # 解析并排序实际输出和预期输出
                    actual_sorted = self._parse_and_sort_output(actual_output)
                    expected_sorted = self._parse_and_sort_output(expected_output)

                    # 若解析成功，则用排序后的值比对
                    if actual_sorted is not None and expected_sorted is not None:
                        actual_output = actual_sorted
                        expected_output = expected_sorted

                # 比对处理后的结果
                if actual_output != expected_output:
                    run_status = "wrong_answer"
                    error = (
                        f"输出不匹配\n"
                        f"预期输出: {test_case['expected_output']}\n"
                        f"实际输出: {output.strip()}"
                    )

            # 记录测试用例结果
            test_case_result = {
                "input": test_case["input_data"],
                "expected_output": test_case["expected_output"],
                "actual_output": output,
                "status": run_status,
                "execution_time": execution_time,
                "error": error,
                "full_code": full_code
            }
            results.append(test_case_result)
            total_time += execution_time

            # 更新整体状态
            if run_status != "accepted" and overall_status == "accepted":
                overall_status = run_status

        return {
            "status": overall_status,
            "results": results,
            "execution_time": total_time,
            "full_code": full_code
        }